/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.protobuf;



import com.google.protobuf.ExtensionRegistry;

import com.google.protobuf.Message;

import com.google.protobuf.Parser;

import com.bff.gaia.unified.sdk.coders.*;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.TypeDescriptor;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableSet;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Sets;

import com.bff.gaia.unified.sdk.coders.CannotProvideCoderException;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderException;

import com.bff.gaia.unified.sdk.coders.CoderProvider;

import com.bff.gaia.unified.sdk.coders.CoderRegistry;

import com.bff.gaia.unified.sdk.coders.CustomCoder;

import com.bff.gaia.unified.sdk.coders.DefaultCoder;



import java.io.IOException;

import java.io.InputStream;

import java.io.OutputStream;

import java.lang.reflect.InvocationTargetException;

import java.lang.reflect.Method;

import java.lang.reflect.Modifier;

import java.util.*;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/**

 * A {@link Coder} using Google Protocol Buffers binary format. {@link ProtoCoder} supports both

 * Protocol Buffers syntax versions 2 and 3.

 *

 * <p>To learn more about Protocol Buffers, visit: <a

 * href="https://developers.google.com/protocol-buffers">https://developers.google.com/protocol-buffers</a>

 *

 * <p>{@link ProtoCoder} is registered in the global {@link CoderRegistry} as the default {@link

 * Coder} for any {@link Message} object. Custom message extensions are also supported, but these

 * extensions must be registered for a particular {@link ProtoCoder} instance and that instance must

 * be registered on the {@link PCollection} that needs the extensions:

 *

 * <pre>{@code

 * import MyProtoFile;

 * import MyProtoFile.MyMessage;

 *

 * Coder<MyMessage> coder = ProtoCoder.of(MyMessage.class).withExtensionsFrom(MyProtoFile.class);

 * PCollection<MyMessage> records = input.apply(...).setCoder(coder);

 * }</pre>

 *

 * <h3>Versioning</h3>

 *

 * <p>{@link ProtoCoder} supports both versions 2 and 3 of the Protocol Buffers syntax. However, the

 * Java runtime version of the <code>google.com.protobuf</code> library must match exactly the

 * version of <code>protoc</code> that was used to produce the JAR files containing the compiled

 * <code>.proto</code> messages.

 *

 * <p>For more information, see the <a

 * href="https://developers.google.com/protocol-buffers/docs/proto3#using-proto2-message-types">Protocol

 * Buffers documentation</a>.

 *

 * <h3>{@link ProtoCoder} and Determinism</h3>

 *

 * <p>In general, Protocol Buffers messages can be encoded deterministically within a single

 * pipeline as long as:

 *

 * <ul>

 *   <li>The encoded messages (and any transitively linked messages) do not use <code>map</code>

 *       fields.

 *   <li>Every Java VM that encodes or decodes the messages use the same runtime version of the

 *       Protocol Buffers library and the same compiled <code>.proto</code> file JAR.

 * </ul>

 *

 * <h3>{@link ProtoCoder} and Encoding Stability</h3>

 *

 * <p>When changing Protocol Buffers messages, follow the rules in the Protocol Buffers language

 * guides for <a href="https://developers.google.com/protocol-buffers/docs/proto#updating">{@code

 * proto2}</a> and <a

 * href="https://developers.google.com/protocol-buffers/docs/proto3#updating">{@code proto3}</a>

 * syntaxes, depending on your message type. Following these guidelines will ensure that the old

 * encoded data can be read by new versions of the code.

 *

 * <p>Generally, any change to the message type, registered extensions, runtime library, or compiled

 * proto JARs may change the encoding. Thus even if both the original and updated messages can be

 * encoded deterministically within a single job, these deterministic encodings may not be the same

 * across jobs.

 *

 * @param <T> the Protocol Buffers {@link Message} handled by this {@link Coder}.

 */

public class ProtoCoder<T extends Message> extends CustomCoder<T> {



  /** Returns a {@link ProtoCoder} for the given Protocol Buffers {@link Message}. */

  public static <T extends Message> ProtoCoder<T> of(Class<T> protoMessageClass) {

    return new ProtoCoder<>(protoMessageClass, ImmutableSet.of());

  }



  /**

   * Returns a {@link ProtoCoder} for the Protocol Buffers {@link Message} indicated by the given

   * {@link TypeDescriptor}.

   */

  public static <T extends Message> ProtoCoder<T> of(TypeDescriptor<T> protoMessageType) {

    @SuppressWarnings("unchecked")

    Class<T> protoMessageClass = (Class<T>) protoMessageType.getRawType();

    return of(protoMessageClass);

  }



  /**

   * Returns a {@link ProtoCoder} like this one, but with the extensions from the given classes

   * registered.

   *

   * <p>Each of the extension host classes must be an class automatically generated by the Protocol

   * Buffers compiler, {@code protoc}, that contains messages.

   *

   * <p>Does not modify this object.

   */

  public ProtoCoder<T> withExtensionsFrom(Iterable<Class<?>> moreExtensionHosts) {

    for (Class<?> extensionHost : moreExtensionHosts) {

      // Attempt to access the required method, to make sure it's present.

      try {

        Method registerAllExtensions =

            extensionHost.getDeclaredMethod("registerAllExtensions", ExtensionRegistry.class);

        checkArgument(

            Modifier.isStatic(registerAllExtensions.getModifiers()),

            "Method registerAllExtensions() must be static");

      } catch (NoSuchMethodException | SecurityException e) {

        throw new IllegalArgumentException(

            String.format("Unable to register extensions for %s", extensionHost.getCanonicalName()),

            e);

      }

    }



    return new ProtoCoder<>(

        protoMessageClass,

        new ImmutableSet.Builder<Class<?>>()

            .addAll(extensionHostClasses)

            .addAll(moreExtensionHosts)

            .build());

  }



  /**

   * See {@link #withExtensionsFrom(Iterable)}.

   *

   * <p>Does not modify this object.

   */

  public ProtoCoder<T> withExtensionsFrom(Class<?>... moreExtensionHosts) {

    return withExtensionsFrom(Arrays.asList(moreExtensionHosts));

  }



  @Override

  public void encode(T value, OutputStream outStream) throws IOException {

    encode(value, outStream, Context.NESTED);

  }



  @Override

  public void encode(T value, OutputStream outStream, Context context) throws IOException {

    if (value == null) {

      throw new CoderException("cannot encode a null " + protoMessageClass.getSimpleName());

    }

    if (context.isWholeStream) {

      value.writeTo(outStream);

    } else {

      value.writeDelimitedTo(outStream);

    }

  }



  @Override

  public T decode(InputStream inStream) throws IOException {

    return decode(inStream, Context.NESTED);

  }



  @Override

  public T decode(InputStream inStream, Context context) throws IOException {

    if (context.isWholeStream) {

      return getParser().parseFrom(inStream, getExtensionRegistry());

    } else {

      return getParser().parseDelimitedFrom(inStream, getExtensionRegistry());

    }

  }



  @Override

  public boolean equals(Object other) {

    if (this == other) {

      return true;

    }

    if (!(other instanceof ProtoCoder)) {

      return false;

    }

    ProtoCoder<?> otherCoder = (ProtoCoder<?>) other;

    return protoMessageClass.equals(otherCoder.protoMessageClass)

        && Sets.newHashSet(extensionHostClasses)

            .equals(Sets.newHashSet(otherCoder.extensionHostClasses));

  }



  @Override

  public int hashCode() {

    return Objects.hash(protoMessageClass, extensionHostClasses);

  }



  @Override

  public void verifyDeterministic() throws NonDeterministicException {

    ProtobufUtil.verifyDeterministic(this);

  }



  /** Returns the Protocol Buffers {@link Message} type this {@link ProtoCoder} supports. */

  public Class<T> getMessageType() {

    return protoMessageClass;

  }



  public Set<Class<?>> getExtensionHosts() {

    return extensionHostClasses;

  }



  /**

   * Returns the {@link ExtensionRegistry} listing all known Protocol Buffers extension messages to

   * {@code T} registered with this {@link ProtoCoder}.

   */

  public ExtensionRegistry getExtensionRegistry() {

    if (memoizedExtensionRegistry == null) {

      ExtensionRegistry registry = ExtensionRegistry.newInstance();

      for (Class<?> extensionHost : extensionHostClasses) {

        try {

          extensionHost

              .getDeclaredMethod("registerAllExtensions", ExtensionRegistry.class)

              .invoke(null, registry);

        } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {

          throw new IllegalStateException(e);

        }

      }

      memoizedExtensionRegistry = registry.getUnmodifiable();

    }

    return memoizedExtensionRegistry;

  }



  ////////////////////////////////////////////////////////////////////////////////////

  // Private implementation details below.



  /** The {@link Message} type to be coded. */

  private final Class<T> protoMessageClass;



  /**

   * All extension host classes included in this {@link ProtoCoder}. The extensions from these

   * classes will be included in the {@link ExtensionRegistry} used during encoding and decoding.

   */

  private final Set<Class<?>> extensionHostClasses;



  // Constants used to serialize and deserialize

  private static final String PROTO_MESSAGE_CLASS = "proto_message_class";

  private static final String PROTO_EXTENSION_HOSTS = "proto_extension_hosts";



  // Transient fields that are lazy initialized and then memoized.

  private transient ExtensionRegistry memoizedExtensionRegistry;

  private transient Parser<T> memoizedParser;



  /** Private constructor. */

  private ProtoCoder(Class<T> protoMessageClass, Set<Class<?>> extensionHostClasses) {

    this.protoMessageClass = protoMessageClass;

    this.extensionHostClasses = extensionHostClasses;

  }



  /** Get the memoized {@link Parser}, possibly initializing it lazily. */

  private Parser<T> getParser() {

    if (memoizedParser == null) {

      try {

        @SuppressWarnings("unchecked")

        T protoMessageInstance = (T) protoMessageClass.getMethod("getDefaultInstance").invoke(null);

        @SuppressWarnings("unchecked")

        Parser<T> tParser = (Parser<T>) protoMessageInstance.getParserForType();

        memoizedParser = tParser;

      } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {

        throw new IllegalArgumentException(e);

      }

    }

    return memoizedParser;

  }



  /**

   * Returns a {@link CoderProvider} which uses the {@link ProtoCoder} for {@link Message proto

   * messages}.

   *

   * <p>This method is invoked reflectively from {@link DefaultCoder}.

   */

  public static CoderProvider getCoderProvider() {

    return new ProtoCoderProvider();

  }



  static final TypeDescriptor<Message> MESSAGE_TYPE = new TypeDescriptor<Message>() {};



  /** A {@link CoderProvider} for {@link Message proto messages}. */

  private static class ProtoCoderProvider extends CoderProvider {



    @Override

    public <T> Coder<T> coderFor(

        TypeDescriptor<T> typeDescriptor, List<? extends Coder<?>> componentCoders)

        throws CannotProvideCoderException {

      if (!typeDescriptor.isSubtypeOf(MESSAGE_TYPE)) {

        throw new CannotProvideCoderException(

            String.format(

                "Cannot provide %s because %s is not a subclass of %s",

                ProtoCoder.class.getSimpleName(), typeDescriptor, Message.class.getName()));

      }



      @SuppressWarnings("unchecked")

      TypeDescriptor<? extends Message> messageType =

          (TypeDescriptor<? extends Message>) typeDescriptor;

      try {

        @SuppressWarnings("unchecked")

        Coder<T> coder = (Coder<T>) ProtoCoder.of(messageType);

        return coder;

      } catch (IllegalArgumentException e) {

        throw new CannotProvideCoderException(e);

      }

    }

  }



  private SortedSet<String> getSortedExtensionClasses() {

    SortedSet<String> ret = new TreeSet<>();

    for (Class<?> clazz : extensionHostClasses) {

      ret.add(clazz.getName());

    }

    return ret;

  }

}